#include <linux/init.h>
#include <linux/module.h>
#include <linux/kthread.h>
#include <linux/completion.h>
#include <linux/timer.h>
#include <linux/sched.h>
#include <linux/delay.h>

#include "speed_def.h"
#include "oppool.h"
#include "debug_file.h"

#define get_driver_name(tfm_type, tfm) crypto_tfm_alg_driver_name(tfm_type##_tfm(tfm))

struct crypto_ahash *hash_tfm;
struct crypto_skcipher *skcipher;

static bool is_running = true;

static struct test_stats test_stats;

static struct timer_list exit_timer;

static struct test_params test_params = { .alg_name = "sm3",
					  .is_hash = true,
					  .num_threads = 1,
					  .test_time_ms = DEFAULT_TEST_TIME_MS,
					  .data_size = DEFAULT_DATA_SIZE,
					  .total_ops = 0,
					  .key_size = 0,
					  .iv_size = 0 };

module_param_string(alg_name, test_params.alg_name, MAX_ALG_NAME_LEN, 0644);
MODULE_PARM_DESC(alg_name,
		 "Algorithm name (e.g., 'cbc(aes)', 'sha256'), default:cbc(aes), default sm3");

module_param_named(is_hash, test_params.is_hash, bool, 0644);
MODULE_PARM_DESC(is_hash,
		 "Set to 1 for hash algorithms, 0 for symmetric crypto algorithms, default:1");

module_param_named(num_threads, test_params.num_threads, int, 0644);
MODULE_PARM_DESC(num_threads, "Number of concurrent threads (1-32), default:1");

module_param_named(test_time_ms, test_params.test_time_ms, int, 0644);
MODULE_PARM_DESC(test_time_ms,
		 "Test duration in milliseconds, default:3000, valid when total_ops is 0");

module_param_named(key_size, test_params.key_size, int, 0644);
MODULE_PARM_DESC(key_size, "Key size in bytes, default:0, max 32");

module_param_named(iv_size, test_params.iv_size, int, 0644);
MODULE_PARM_DESC(iv_size, "IV size in bytes, default:0, max 16");

module_param_named(total_ops, test_params.total_ops, long, 0644);
MODULE_PARM_DESC(total_ops, "Total operations to perform (0 for time-based), default:0");

void set_datasize(int size)
{
	test_params.data_size = size;
}

static int kthread_worker_count(void *data)
{
	struct thread_data *tdata = (struct thread_data *)data;
	int ret = init_async_op_pool(tdata);
	long send_ops = 0;
	struct async_op *ops[BATCH_SIZE] = { 0 };
	int ops_count = 0, i = 0, j = 0;

	if (ret) {
		pr_err("Failed to initialize async op pool (err: %d)\n", ret);
		return ret;
	}

	// 记录开始时间
	tdata->start_time = ktime_get();
	if (tdata->params->is_hash) {
		while (send_ops < tdata->params->total_ops) {
			ops_count = get_async_ops_batch(tdata, ops, BATCH_SIZE, 1);
			if (ops_count <= 0) {
				continue;
			}

			for (i = 0; i < ops_count; i++) {
				if (ops[i] == NULL) {
					pr_err("get null ops, i: %d\n", i);
					continue;
				}

				crypto_init_wait(&ops[i]->wait);
				ahash_request_set_crypt(ops[i]->ahreq, &ops[i]->sg_src,
							ops[i]->dst_buf, tdata->params->data_size);

				ret = crypto_ahash_digest(ops[i]->ahreq);
				if (ret == 0) {
					/* sync */
					atomic64_inc(&ops[i]->pool->executed_ops);
					put_async_op2pool(ops[i]->pool, ops[i]);
				} else {
					if (ret == -EINPROGRESS || ret == -EBUSY) {
						// crypto_wait_req(ret, &ops[i]->wait);
						// put_async_op(tdata,
						// 	     ops[i]); /* TODO put in callback ??*/

						// pr_info("op %p, src %p, dst %p, len %d i %d\n",
						// 	ops[i], ops[i]->src_buf, ops[i]->dst_buf,
						// 	tdata->params->data_size, i);
					} else {
						/* error */
						pr_err("Failed to encrypt async op (err: %d)\n",
						       ret);
						/* put back to pool */
						for (j = i; j < ops_count; j++)
							put_async_op2pool(ops[i]->pool, ops[j]);

						break;
					}
				}
			}
			send_ops += i;
		}
	} else {
		while (send_ops < tdata->params->total_ops) {
			ops_count = get_async_ops_batch(tdata, ops, BATCH_SIZE, 1);
			if (ops_count <= 0) {
				continue;
			}

			for (i = 0; i < ops_count; i++) {
				if (ops[i] == NULL) {
					pr_err("get null ops, i: %d\n", i);
					continue;
				}

				crypto_init_wait(&ops[i]->wait);
				// skcipher_request_set_callback(ops[i]->skreq,
				// 			      CRYPTO_TFM_REQ_MAY_BACKLOG,
				// 			      async_skcipher_op_complete, ops[i]);
				skcipher_request_set_crypt(ops[i]->skreq, &ops[i]->sg_src,
							   &ops[i]->sg_dst,
							   tdata->params->data_size,
							   tdata->params->iv);
				ret = crypto_skcipher_encrypt(ops[i]->skreq);
				if (ret == 0) {
					/* sync */
					atomic64_inc(&ops[i]->pool->executed_ops);
					put_async_op2pool(ops[i]->pool, ops[i]);
				} else {
					if (ret == -EINPROGRESS || ret == -EBUSY) {
						// crypto_wait_req(ret, &ops[i]->wait);
						// put_async_op(tdata, ops[i]);
						// pr_info("async process i %d done\n", i);
					} else {
						/* error */
						pr_err("Failed to encrypt async op (err: %d)\n",
						       ret);
						/* put back to pool */
						for (j = i; j < ops_count; j++)
							put_async_op2pool(ops[i]->pool, ops[j]);

						break;
					}
				}

				// pr_info("ret %d, ops_count %d i %d\n", ret, ops_count, i);
			}
			send_ops += i;
			// pr_info("total %d, ops_count %d\n", total, ops_count);
		}
	}

	i = 0;
	while (atomic64_read(&tdata->op_pool->executed_ops) < send_ops) {
		schedule_timeout_interruptible(msecs_to_jiffies(1));
		i++;
		if (i > 10) { /* after 10ms */
			pr_err("Failed to wait for async ops (executed_ops: %llu), total: %lu\n",
			       atomic64_read(&tdata->op_pool->executed_ops), send_ops);
			break;
		}
	}

	send_ops = atomic64_read(&tdata->op_pool->executed_ops);

	// 记录结束时间
	tdata->end_time = ktime_get();

	// 计算速度（Mbps）
	tdata->speed_mbps = (u64)(tdata->params->data_size * 8 * send_ops * 1000) /
			    ktime_to_ns(tdata->end_time - tdata->start_time);
	pr_info("Thread[%d]:[%d bytes] %llu Mbps done ops %lu\n", tdata->thread_id,
		tdata->params->data_size, tdata->speed_mbps, send_ops);

	// 清理异步操作池
	cleanup_async_op_pool(tdata);

	// 线程退出前，唤醒等待者
	complete(&tdata->completion);

	return 0;
}

// 内核线程函数
static int kthread_worker_timer(void *data)
{
	struct thread_data *tdata = (struct thread_data *)data;
	int ret = init_async_op_pool(tdata);
	u64 total = 0;
	struct async_op *ops[BATCH_SIZE] = { 0 };
	int ops_count = 0, i = 0, j = 0;

	if (ret) {
		pr_err("Failed to initialize async op pool (err: %d)\n", ret);
		return ret;
	}

	// 记录开始时间
	tdata->start_time = ktime_get();
	if (tdata->params->is_hash) {
		while (is_running) {
			ops_count = get_async_ops_batch(tdata, ops, BATCH_SIZE, 1);
			if (ops_count <= 0) {
				continue;
			}

			for (i = 0; i < ops_count; i++) {
				if (ops[i] == NULL) {
					pr_err("get null ops, i: %d\n", i);
					continue;
				}

				crypto_init_wait(&ops[i]->wait);
				ahash_request_set_crypt(ops[i]->ahreq, &ops[i]->sg_src,
							ops[i]->dst_buf, tdata->params->data_size);

				ret = crypto_ahash_digest(ops[i]->ahreq);
				if (ret == 0) {
					/* sync */
					atomic64_inc(&ops[i]->pool->executed_ops);
					put_async_op(tdata, ops[i]);
				} else {
					if (ret == -EINPROGRESS || ret == -EBUSY) {
						// crypto_wait_req(ret, &ops[i]->wait);
						// put_async_op(tdata,
						// 	     ops[i]); /* TODO put in callback ??*/

						// pr_info("op %p, src %p, dst %p, len %d i %d\n",
						// 	ops[i], ops[i]->src_buf, ops[i]->dst_buf,
						// 	tdata->params->data_size, i);
					} else {
						/* error */
						pr_err("Failed to encrypt async op (err: %d)\n",
						       ret);
						/* put back to pool */
						for (j = i; j < ops_count; j++)
							put_async_op(tdata, ops[j]);

						break;
					}
				}
			}
			total += i;
		}
	} else {
		while (is_running) {
			ops_count = get_async_ops_batch(tdata, ops, BATCH_SIZE, 1);
			if (ops_count <= 0) {
				continue;
			}

			for (i = 0; i < ops_count; i++) {
				if (ops[i] == NULL) {
					pr_err("get null ops, i: %d\n", i);
					continue;
				}

				crypto_init_wait(&ops[i]->wait);
				// skcipher_request_set_callback(ops[i]->skreq,
				// 			      CRYPTO_TFM_REQ_MAY_BACKLOG,
				// 			      async_skcipher_op_complete, ops[i]);
				skcipher_request_set_crypt(ops[i]->skreq, &ops[i]->sg_src,
							   &ops[i]->sg_dst,
							   tdata->params->data_size,
							   tdata->params->iv);
				ret = crypto_skcipher_encrypt(ops[i]->skreq);
				if (ret == 0) {
					/* sync */
					atomic64_inc(&ops[i]->pool->executed_ops);
					put_async_op(tdata, ops[i]);
				} else {
					if (ret == -EINPROGRESS || ret == -EBUSY) {
						// crypto_wait_req(ret, &ops[i]->wait);
						// put_async_op(tdata, ops[i]);
						// pr_info("async process i %d done\n", i);
					} else {
						/* error */
						pr_err("Failed to encrypt async op (err: %d)\n",
						       ret);
						/* put back to pool */
						for (j = i; j < ops_count; j++)
							put_async_op(tdata, ops[j]);

						break;
					}
				}

				// pr_info("ret %d, ops_count %d i %d\n", ret, ops_count, i);
			}
			total += i;
			// pr_info("total %d, ops_count %d\n", total, ops_count);
		}
	}

	i = 0;
	while (atomic64_read(&tdata->op_pool->executed_ops) < total) {
		schedule_timeout_interruptible(msecs_to_jiffies(1));
		i++;
		if (i > 10) { /* after 10ms */
			pr_err("Failed to wait for async ops (executed_ops: %llu), total: %llu\n",
			       atomic64_read(&tdata->op_pool->executed_ops), total);
			break;
		}
	}

	total = atomic64_read(&tdata->op_pool->executed_ops);

	// 记录结束时间
	tdata->end_time = ktime_get();

	// 计算速度（Mbps）
	tdata->speed_mbps = (u64)(tdata->params->data_size * 8 * total * 1000) /
			    ktime_to_ns(tdata->end_time - tdata->start_time);
	pr_info("Thread[%d]:[%d bytes] %llu Mbps done ops %llu\n", tdata->thread_id,
		tdata->params->data_size, tdata->speed_mbps, total);

	// 清理异步操作池
	cleanup_async_op_pool(tdata);

	// 线程退出前，唤醒等待者
	complete(&tdata->completion);

	return 0;
}

static void timer_handler(struct timer_list *t)
{
	pr_info("Timer triggered: stopping all threads\n");
	is_running = false; // 所有线程会检测到该标志并退出
}

static void start_performance_test(void)
{
	int ret, i;
	u64 speed_mbps = 0;

	/* alloc tfm first */
	if (test_params.is_hash) {
		hash_tfm = crypto_alloc_ahash(test_params.alg_name, 0, 0);
		if (IS_ERR(hash_tfm)) {
			ret = PTR_ERR(hash_tfm);
			pr_err("Failed to allocate hash tfm (err: %d)\n", ret);
			return;
		}
		pr_info("Get %s driver: %s\n", test_params.alg_name,
			get_driver_name(crypto_ahash, hash_tfm));

		if (test_params.key_size > 0) {
			ret = crypto_ahash_setkey(hash_tfm, test_params.key, test_params.key_size);
			if (ret) {
				crypto_free_ahash(hash_tfm);
				pr_err("Failed to set hash key (err: %d)\n", ret);
				return;
			}
		}
	} else {
		skcipher = crypto_alloc_skcipher(test_params.alg_name, 0, 0);
		if (IS_ERR(skcipher)) {
			ret = PTR_ERR(skcipher);
			pr_err("Failed to allocate skcipher tfm (err: %d)\n", ret);
			return;
		}

		pr_info("%s driver: %s\n", test_params.alg_name,
			get_driver_name(crypto_skcipher, skcipher));

		ret = crypto_skcipher_setkey(skcipher, test_params.key, test_params.key_size);
		if (ret) {
			crypto_free_skcipher(skcipher);
			pr_err("Failed to set skcipher key (err: %d)\n", ret);
			return;
		}
	}

	is_running = true;
	// 初始化所有线程数据并创建线程
	for (i = 0; i < test_params.num_threads; i++) {
		// 初始化线程数据
		test_stats.threads[i].thread_id = i;
		init_completion(&test_stats.threads[i].completion);
		test_stats.threads[i].params = &test_params;

		if (test_params.total_ops == 0) {
			// 创建定时内核线程
			test_stats.threads[i].thread =
				kthread_run(kthread_worker_timer, // 线程函数
					    &test_stats.threads[i], // 传递当前线程数据作为参数
					    "kthread-%d", i // 线程名称
				);
		} else {
			test_stats.threads[i].thread =
				kthread_run(kthread_worker_count, // 线程函数
					    &test_stats.threads[i], // 传递当前线程数据作为参数
					    "kthread-%d", i // 线程名称
				);
		}

		// 检查线程创建是否成功
		if (IS_ERR(test_stats.threads[i].thread)) {
			ret = PTR_ERR(test_stats.threads[i].thread);
			pr_err("Failed to create thread %d (err: %d)\n", i, ret);
			goto cleanup_threads; // 清理已创建的线程
		}
	}

	if (test_params.total_ops == 0) {
		timer_setup(&exit_timer, timer_handler, 0); // 第三个参数为0表示不使用数据传递
		exit_timer.expires = jiffies + msecs_to_jiffies(test_params.test_time_ms);
		add_timer(&exit_timer);
		pr_info("Timer started (%d ms delay)\n", test_params.test_time_ms);
	}

	// 等待所有线程退出
	for (i = 0; i < test_params.num_threads; i++) {
		wait_for_completion(&test_stats.threads[i].completion);
		reinit_completion(&test_stats.threads[i].completion);
		pr_info("Thread %d exited successfully\n", i);
		speed_mbps += test_stats.threads[i].speed_mbps;
		test_stats.threads[i].thread = NULL;
	}

	// 打印所有线程的速度
	pr_info("Summary[%s : %d bytes] %llu Mbps\n", test_params.alg_name, test_params.data_size,
		speed_mbps);
	pr_info("==================================\n");

	if (test_params.total_ops == 0)
		del_timer_sync(&exit_timer);

	if (hash_tfm)
		crypto_free_ahash(hash_tfm);

	if (skcipher)
		crypto_free_skcipher(skcipher);

	return;

cleanup_threads:

	is_running = false;

	// 清理所有线程
	for (i = 0; i < test_params.num_threads; i++) {
		if (test_stats.threads[i].thread) {
			// kthread_stop(test_stats.threads[i].thread);
			pr_info("Thread %d stopped\n", i);
		}
	}

	return;
}

// Tasklet函数 - 执行性能测试
void performance_test_tasklet(unsigned long data)
{
	start_performance_test();
}

static int __init my_module_init(void)
{
	int ret = 0;

	// 校验线程数量合法性
	if (test_params.num_threads <= 0 || test_params.num_threads > MAX_THREADS) {
		pr_err("Invalid num_threads: %d (must be 1-%d)\n", test_params.num_threads,
		       MAX_THREADS);
		return -EINVAL;
	}

	if (test_params.key_size > 0) {
		if (test_params.key_size > MAX_KEY_SIZE) {
			pr_err("Invalid key_size: %d (must be 0-%d)\n", test_params.key_size,
			       MAX_KEY_SIZE);
			return -EINVAL;
		}
		get_random_bytes(test_params.key, test_params.key_size);
	}

	if (test_params.iv_size > 0) {
		if (test_params.iv_size > MAX_IV_SIZE) {
			pr_err("Invalid iv_size: %d (must be 0-%d)\n", test_params.iv_size,
			       MAX_IV_SIZE);
			return -EINVAL;
		}
		get_random_bytes(test_params.iv, test_params.iv_size);
	}

	// 初始化debugfs文件
	ret = debug_file_init();
	if (ret) {
		pr_err("Failed to initialize debugfs files (err: %d)\n", ret);
		return ret;
	}

	pr_info("Get params: %s, num_threads %d, data_size %d, test_time_ms %d, total_ops %ld, is_hash %d\n",
		test_params.alg_name, test_params.num_threads, test_params.data_size,
		test_params.test_time_ms, test_params.total_ops, test_params.is_hash);
	pr_info("Load kspeed module success!\n");

	return 0;
}

static void __exit my_module_exit(void)
{
	// 清理debugfs文件
	debug_file_exit();

	pr_info("Module exit: all resources cleaned up\n");
}

module_init(my_module_init);
module_exit(my_module_exit);
MODULE_LICENSE("GPL");
